package cover

import (
	"container/heap"
	"github.com/viant/vec/search"
	"io"
	"math"
	"sort"
	"sync"
)

// Tree represents a cover tree.
type Tree[T any] struct {
	root             *Node
	base             float32
	distanceFuncName DistanceFunction
	distanceFnunc    DistanceFunc
	values           values[T]
	indexMap         map[int32]*Point
	version          uint64
	boundStrategy    BoundStrategy
	mu               sync.RWMutex
}

// BoundStrategy selects which lower-bound radius to use when pruning.
type BoundStrategy int

const (
	// BoundPerNode uses cached per-node subtree radius (tighter, more pruning).
	BoundPerNode BoundStrategy = iota
	// BoundLevel uses conservative geometric bound from node.baseLevel (pow-free).
	BoundLevel
)

// Insert adds a new point (embedding vector) to the cover tree.
func (t *Tree[T]) Insert(value T, point *Point) int32 {
	t.mu.Lock()
	defer t.mu.Unlock()
	point.index = t.values.put(value)
	if t.indexMap == nil {
		t.indexMap = make(map[int32]*Point)
	}
	t.indexMap[point.index] = point //
	if point.Magnitude == 0 && len(point.Vector) > 0 {
		point.Magnitude = search.Float32s(point.Vector).Magnitude()
	}
	if t.root == nil {
		t.root = &Node{point: point, level: 0}
	} else {
		t.insert(t.root, point, 0)
	}
	// bump version to invalidate cached radii
	t.version++
	return point.index
}

// FindPointByIndex returns the point associated with the given index.
func (t *Tree[T]) FindPointByIndex(index int32) *Point {
	t.mu.RLock()
	defer t.mu.RUnlock()
	if point, exists := t.indexMap[index]; exists {
		return point
	}
	return nil
}

func (t *Tree[T]) EncodeValues(writer io.Writer) error {
	t.mu.RLock()
	defer t.mu.RUnlock()
	return t.values.Encode(writer)
}

func (t *Tree[T]) DecodeValues(reader io.Reader) error {
	t.mu.Lock()
	defer t.mu.Unlock()
	t.values = values[T]{data: make([]T, 0)}
	t.values.ensureType()
	return t.values.Decode(reader)
}

func (t *Tree[T]) EncodeTree(writer io.Writer) error {
	t.mu.RLock()
	defer t.mu.RUnlock()
	buffer := writers.Get()
	defer writers.Put(buffer)
	buffer.Float32(t.base)

	buffer.String(string(t.distanceFuncName))
	if err := buffer.Coder(t.root); err != nil {
		return err
	}
	_, err := writer.Write(buffer.Bytes())
	return err
}

func (t *Tree[T]) DecodeTree(reader io.Reader) error {
	t.mu.Lock()
	defer t.mu.Unlock()
	buffer := readers.Get()
	defer readers.Put(buffer)
	data, err := io.ReadAll(reader)
	if err != nil {
		return err
	}
	if err = buffer.FromBytes(data); err != nil {
		return err
	}
	buffer.Float32(&t.base)
	var distance string
	buffer.String(&distance)
	t.distanceFuncName = DistanceFunction(distance)
	t.distanceFnunc = t.distanceFuncName.Function()
	t.root = &Node{}
	if err := buffer.Coder(t.root); err != nil {
		return err
	}
	// new structure loaded; bump version to invalidate any cached radii
	t.version++
	return nil
}

func (t *Tree[T]) insert(node *Node, point *Point, level int32) {
	for {
		baseLevel := float32(math.Pow(float64(t.base), float64(level)))
		distance := t.distanceFnunc(point, node.point)

		if distance < baseLevel {
			// Check if the point can be inserted into any child
			inserted := false
			for i := range node.children {
				child := &node.children[i]
				if t.distanceFnunc(point, child.point) < baseLevel {
					node = child
					level--
					inserted = true
					break
				}
			}
			if !inserted {
				// Insert as a new child
				node.children = append(node.children, NewNode(point, level-1, t.base))
				return
			}
		} else {
			// Promote the point to a higher level
			level++
			if level > node.level {
				newRoot := NewNode(point, level, t.base)
				newRoot.children = append(newRoot.children, *t.root)
				t.root = &newRoot
				return
			}
		}
	}
}

// Remove removes a point (embedding vector) from the cover tree.
func (t *Tree[T]) Remove(point *Point) bool {
	t.mu.Lock()
	defer t.mu.Unlock()
	if t.root == nil {
		return false
	}
	removed, newRoot := t.remove(t.root, point)
	t.root = newRoot
	if removed {
		var empty T
		t.values.data[point.index] = empty // Remove the value from the slice
		delete(t.indexMap, point.index)    // Remove the point from the map
		// structure changed; invalidate cached radii
		t.version++
	}
	return removed
}

func (t *Tree[T]) remove(node *Node, point *Point) (bool, *Node) {
	if node == nil {
		return false, nil
	}
	if t.distanceFnunc(point, node.point) == 0 {
		if len(node.children) == 0 {
			return true, nil
		}

		// Promote one of the children to be the new node
		newNode := &node.children[0]
		for _, child := range node.children[1:] {
			t.insert(newNode, child.point, child.level)
		}
		return true, newNode
	}
	for i := range node.children {
		child := &node.children[i]
		removed, newChild := t.remove(child, point)
		if removed {
			if newChild == nil {
				node.children = append(node.children[:i], node.children[i+1:]...)
			} else {
				node.children[i] = *newChild
			}
			return true, node
		}
	}
	return false, node
}

func (t *Tree[T]) Value(point *Point) T {
	t.mu.RLock()
	defer t.mu.RUnlock()
	var r T
	if point == nil || !point.HasValue() {
		return r
	}
	return t.values.value(point.index)
}

func (t *Tree[T]) Values(points []*Point) []T {
	t.mu.RLock()
	defer t.mu.RUnlock()
	var result = make([]T, 0, len(points))
	for i, point := range points {
		if point == nil || point.index < 0 {
			continue
		}
		result[i] = t.values.value(point.index)
	}
	return result
}

// KNearestNeighbors finds the k nearest neighbors of the given point (embedding vector) in the cover tree.
func (t *Tree[T]) KNearestNeighbors(point *Point, k int) []*Neighbor {
	// Lock strategy: per-node radius mutates cache; use write lock. Level-based bound uses read lock.
	if t.boundStrategy == BoundPerNode {
		t.mu.Lock()
		defer t.mu.Unlock()
	} else {
		t.mu.RLock()
		defer t.mu.RUnlock()
	}
	if t.root == nil {
		return nil
	}
	h := &Neighbors{}
	heap.Init(h)
	t.kNearestNeighbors(t.root, point, k, h)
	result := make([]*Neighbor, h.Len())
	for i := len(result) - 1; i >= 0; i-- {
		n := heap.Pop(h).(Neighbor)
		result[i] = &n
	}
	return result
}

// ensureRadius returns the per-node subtree radius, computing and caching it if needed.
// radius(node) = max over children { dist(node, child.point) + radius(child) }.
func (t *Tree[T]) ensureRadius(n *Node) float32 {
	if n == nil {
		return 0
	}
	// If version matches and computed, return cached
	if n.radiusComputed == t.version {
		return n.radius
	}
	if len(n.children) == 0 {
		n.radius = 0
		n.radiusComputed = t.version
		return 0
	}
	maxR := float32(0)
	for i := range n.children {
		c := &n.children[i]
		cr := t.ensureRadius(c)
		d := t.distanceFnunc(n.point, c.point) + cr
		if d > maxR {
			maxR = d
		}
	}
	n.radius = maxR
	n.radiusComputed = t.version
	return maxR
}

func (t *Tree[T]) kNearestNeighbors(node *Node, point *Point, k int, h *Neighbors) {
	// 1) consider node center
	dc := t.distanceFnunc(point, node.point)
	if h.Len() < k {
		heap.Push(h, Neighbor{Point: node.point, Distance: dc})
	} else if dc < (*h)[0].Distance {
		heap.Pop(h)
		heap.Push(h, Neighbor{Point: node.point, Distance: dc})
	}

	if len(node.children) == 0 {
		return
	}

	// 2) compute distances to children
	type childDist struct {
		child *Node
		dist  float32
	}
	cds := make([]childDist, 0, len(node.children))
	for i := range node.children {
		c := &node.children[i]
		cds = append(cds, childDist{child: c, dist: t.distanceFnunc(point, c.point)})
	}

	// 3) process children closest-first
	sort.Slice(cds, func(i, j int) bool { return cds[i].dist < cds[j].dist })

	// 4) branch-and-bound using configured radius strategy
	for _, cd := range cds {
		var worstBest float32
		if h.Len() == k {
			worstBest = (*h)[0].Distance
		} else {
			worstBest = float32(math.MaxFloat32)
		}
		r := t.boundRadius(cd.child)
		lb := cd.dist - r
		if h.Len() == k && lb >= worstBest {
			// prune this subtree
			continue
		}
		t.kNearestNeighbors(cd.child, point, k, h)
	}
}

// KNearestNeighborsBestFirst finds k nearest neighbors using a best-first (node-PQ) search.
// It maintains a priority queue of subtrees keyed by a conservative lower bound and
// pops the smallest lower bound first to maximize pruning effectiveness.
func (t *Tree[T]) KNearestNeighborsBestFirst(point *Point, k int) []*Neighbor {
	if t.boundStrategy == BoundPerNode {
		t.mu.Lock()
		defer t.mu.Unlock()
	} else {
		t.mu.RLock()
		defer t.mu.RUnlock()
	}
	if t.root == nil {
		return nil
	}
	// Max-heap of best neighbors
	nh := &Neighbors{}
	heap.Init(nh)

	// Min-heap of nodes ordered by lower bound
	pq := &nodeQueue{}
	heap.Init(pq)

	// Seed with root
	rootDist := t.distanceFnunc(point, t.root.point)
	rootLB := rootDist - t.boundRadius(t.root)
	heap.Push(pq, nodeItem{node: t.root, lb: rootLB, centerDist: rootDist})

	for pq.Len() > 0 {
		worstBest := float32(math.MaxFloat32)
		if nh.Len() == k {
			worstBest = (*nh)[0].Distance
		}
		top := heap.Pop(pq).(nodeItem)
		if nh.Len() == k && top.lb >= worstBest {
			// All remaining nodes have lb >= top.lb >= worstBest: early stop
			break
		}

		// Consider the center of this node
		dc := top.centerDist
		if nh.Len() < k {
			heap.Push(nh, Neighbor{Point: top.node.point, Distance: dc})
		} else if dc < (*nh)[0].Distance {
			heap.Pop(nh)
			heap.Push(nh, Neighbor{Point: top.node.point, Distance: dc})
		}

		// Enqueue children with their lower bounds
		for i := range top.node.children {
			child := &top.node.children[i]
			cd := t.distanceFnunc(point, child.point)
			lb := cd - t.boundRadius(child)
			if nh.Len() == k && lb >= (*nh)[0].Distance {
				continue
			}
			heap.Push(pq, nodeItem{node: child, lb: lb, centerDist: cd})
		}
	}

	// Drain neighbor heap into ascending order
	result := make([]*Neighbor, nh.Len())
	for i := len(result) - 1; i >= 0; i-- {
		n := heap.Pop(nh).(Neighbor)
		result[i] = &n
	}
	return result
}

// nodeCoverRadius returns the per-node cached subtree radius (computing if needed).
func (t *Tree[T]) nodeCoverRadius(n *Node) float32 { return t.ensureRadius(n) }

// levelCoverRadius returns the conservative geometric bound R(L) = base^(L+1)/(base-1)
// computed without pow by using node.baseLevel = base^L.
func (t *Tree[T]) levelCoverRadius(n *Node) float32 {
	if t.base <= 1 || n == nil {
		return float32(math.MaxFloat32)
	}
	return n.baseLevel * t.base / (t.base - 1)
}

// boundRadius selects the radius based on the configured strategy.
func (t *Tree[T]) boundRadius(n *Node) float32 {
	if t.boundStrategy == BoundLevel {
		return t.levelCoverRadius(n)
	}
	return t.nodeCoverRadius(n)
}

// SetBoundStrategy updates the bound strategy at runtime.
func (t *Tree[T]) SetBoundStrategy(s BoundStrategy) { t.boundStrategy = s }

// nodeItem represents a subtree with a lower bound and cached center distance
type nodeItem struct {
	node       *Node
	lb         float32
	centerDist float32
}

// nodeQueue is a min-heap by lower bound
type nodeQueue []nodeItem

func (q nodeQueue) Len() int            { return len(q) }
func (q nodeQueue) Less(i, j int) bool  { return q[i].lb < q[j].lb }
func (q nodeQueue) Swap(i, j int)       { q[i], q[j] = q[j], q[i] }
func (q *nodeQueue) Push(x interface{}) { *q = append(*q, x.(nodeItem)) }
func (q *nodeQueue) Pop() interface{} {
	old := *q
	n := len(old)
	x := old[n-1]
	*q = old[:n-1]
	return x
}

// NewTree initializes and returns a new Tree.
func NewTree[T any](base float32, distanceFn DistanceFunction) *Tree[T] {
	return &Tree[T]{base: base, distanceFnunc: distanceFn.Function(), distanceFuncName: distanceFn, values: values[T]{data: make([]T, 0)}, boundStrategy: BoundPerNode}
}
